﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using TorchSharp;
using static TorchSharp.torch;

namespace NgpSharp.RayMaching
{
    /// <summary>
    /// 
    /// </summary>
    public class Morton3DFounction : autograd.SingleTensorFunction<Morton3DFounction>
    {
        /// <summary>
        /// 
        /// </summary>
        public override string Name => nameof(Morton3DFounction);

        public override Tensor forward(autograd.AutogradContext ctx, params object[] vars)
        {
            using var c = torch.NewDisposeScope();
            Tensor coords = (Tensor)vars[0];
            long N = coords.shape[0];
            if(!coords.is_cuda) coords = coords.cuda();
            Tensor indices = torch.empty(N, dtype: torch.int32, device: coords.device);
            coords = coords.to(ScalarType.Int32);
            RayMachingLoader.morton3D(coords.Handle, (uint)N, indices.Handle);
            return indices.MoveToOuterDisposeScope();
        }

        public override List<Tensor> backward(autograd.AutogradContext ctx, Tensor grad_output)
        {
            grad_output.Dispose();
            return [null];
        }
    }
}
